#define K_MAX_SHAPE_DIM 0
#include "kernel_operator.h"
#include <type_traits>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
template <typename T>
class KernelReshape_all {
public:
    __aicore__ inline KernelReshape_all() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y,
        uint32_t CoreDataNum, uint32_t finalTileNum, uint32_t tileDataNum, uint32_t TailDataNum, TPipe* pipe) {

        this->coreDataNum = CoreDataNum;
        this->tileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->tailDataNum = TailDataNum;
        
        xGm.SetGlobalBuffer((__gm__ T*)x, this->coreDataNum);
        yGm.SetGlobalBuffer((__gm__ T*)y, this->coreDataNum);

        pipe->InitBuffer(QueueTmp1, this->tileDataNum * sizeof(T));
    }
    __aicore__ inline void Process() {
        int32_t loopCount = this->tileNum;
        uint32_t currOffset = 0;
        auto tmp1 = QueueTmp1.Get<T>();
        for (int32_t i = 0; i < loopCount; i++) {
            DataCopy(tmp1, xGm[currOffset], this->tileDataNum);
            DataCopy(yGm[currOffset], tmp1, this->tileDataNum);
            currOffset += this->tileDataNum;
        }
        DataCopy(tmp1, xGm[currOffset], this->tailDataNum);
        DataCopy(yGm[currOffset], tmp1, this->tailDataNum);
    }
private:
    TPipe* pipe;
    TQueBind<QuePosition::VECIN, QuePosition::VECOUT, BUFFER_NUM> inOutQueue; 
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2,QueueTmp3,QueueTmp4, QueueTmp5;

    GlobalTensor<T> xGm;
    GlobalTensor<T> yGm;
    
    uint32_t coreDataNum;
    uint32_t tileNum;
    uint32_t tileDataNum;
    uint32_t tailDataNum;

    uint32_t processDataNum;
};
extern "C" __global__ __aicore__ void reshape(GM_ADDR x, GM_ADDR shape, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelReshape_all<DTYPE_X> op;
    TPipe pipe;
    op.Init(x, y, 
        tiling_data.CoreDataNum, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum, &pipe);
    op.Process();
}